package com.helmet_back.controller;

import com.helmet_back.utils.Command;
import com.helmet_back.utils.FilePath;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;

import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.nio.file.StandardCopyOption;
import java.util.Base64;


@RestController
@RequestMapping("/yolo")
public class ImgController {



    @PostMapping("/detect")
    public String receive(@RequestParam("source") MultipartFile file){
//        long start = System.currentTimeMillis(); // 记录开始时间
        String type="";
        if (file.getContentType().startsWith("image")) {
            // 上传的文件是图片
            type="image";
        } else if (file.getContentType().startsWith("video")) {
            // 上传的文件是视频
            type = "video";
        }
        try {

            Path uploadPath = Paths.get(FilePath.uploadPath);
            if (!Files.exists(uploadPath)) {
                Files.createDirectories(uploadPath);
            }
            Files.copy(file.getInputStream(),
                    uploadPath.resolve(file.getOriginalFilename()),
                    StandardCopyOption.REPLACE_EXISTING);

            String command = "python detect.py --source ./sources/"+file.getOriginalFilename()+" --weight ./weights/bests.pt"; // 要执行的cmd命令
            String workingDirectory = FilePath.yoloPath; // cmd的工作目录
//            long start = System.currentTimeMillis(); // 记录开始时间
            String result = Command.executeCmd(command,workingDirectory);
            System.out.println(result);
            Path uploadedFilePath = uploadPath.resolve(file.getOriginalFilename());
            Files.deleteIfExists(uploadedFilePath);

//            long end = System.currentTimeMillis(); // 记录结束时间
//            long executionTime = end - start; // 计算执行时间
//            System.out.println("springboot执行时间为：" + executionTime/1000 + " 秒");

            if(type.equals("image")){
                return returnImgURL(file);
            } else if (type.equals("video")) {
                return returnVideoURL(file);
            }
        }catch (Exception e){
            System.out.println(e);
        }
        return "error!";

    }


    public String returnImgURL(@RequestParam("files") MultipartFile file) throws IOException {
        String base64Image = null;
        try {
            String imagePath = FilePath.outputPath + file.getOriginalFilename(); // 本地图片路径
            File imageFile = new File(imagePath);
            byte[] imageBytes = Files.readAllBytes(imageFile.toPath());
            base64Image = Base64.getEncoder().encodeToString(imageBytes);
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            return "data:image/png;base64,"+ base64Image;
        }
    }

//    public String returnVideoURL(@RequestParam("files") MultipartFile file) throws IOException {
//        String base64Image = null;
//        try {
//            System.out.println("video 检测");
//            String imagePath = FilePath.outputPath + file.getOriginalFilename();
//            File imageFile = new File(imagePath);
//            byte[] imageBytes = Files.readAllBytes(imageFile.toPath());
//            base64Image = Base64.getEncoder().encodeToString(imageBytes);
//        } catch (IOException e) {
//            e.printStackTrace();
//        } finally {
//            //return "data:video/mp4;base64,"+ base64Image;
////            System.out.println("返回");
//            return base64Image;
//        }
//
//    }
public String returnVideoURL(@RequestParam("files") MultipartFile file){
    if (!file.isEmpty()) {
        // 定义上传文件的目录
        String uploadDir = "D:\\Code\\python\\helmet\\Smart_Construction\\inference\\output\\"+file.getOriginalFilename();


        return uploadDir;
    } else {
        return "File upload failed!";
    }
    }


}
